import dataclasses
import os
import re
import time
from typing import Any, Dict, List, Tuple
from uuid import uuid1

import loguru
import requests
from tenacity import *

from pentestgpt.utils.llm_api import LLMAPI

logger = loguru.logger
logger.remove()


@dataclasses.dataclass
class Message:
    ask_id: str = None
    ask: dict = None
    answer: dict = None
    answer_id: str = None
    request_start_timestamp: float = None
    request_end_timestamp: float = None
    time_escaped: float = None


@dataclasses.dataclass
class Conversation:
    conversation_id: str = None
    message_list: List[Message] = dataclasses.field(default_factory=list)

    def __hash__(self):
        return hash(self.conversation_id)

    def __eq__(self, other):
        if not isinstance(other, Conversation):
            return False
        return self.conversation_id == other.conversation_id


class OllamaAPI(LLMAPI):
    def __init__(self, config_class, use_langfuse_logging=False):
        self.name = str(config_class.model)
        self.history_length = 10  # maintain conversation history
        self.config = config_class
        self.model = config_class.model
        self.api_base = config_class.api_base
        self.log_dir = config_class.log_dir
        self.conversation_dict: Dict[str, Conversation] = {}
        self.use_langfuse_logging = use_langfuse_logging
        
        if self.log_dir:
            logger.add(sink=os.path.join(self.log_dir, "ollama.log"), level="WARNING")

        # Test connection to Ollama
        self._test_connection()

    def _test_connection(self):
        """Test if Ollama server is running and model is available"""
        try:
            # Check if Ollama server is running
            response = requests.get(f"{self.api_base}/api/tags", timeout=5)
            if response.status_code != 200:
                raise Exception(f"Ollama server not accessible at {self.api_base}")
            
            # Check if model is available
            models = response.json().get("models", [])
            model_names = [model["name"] for model in models]
            
            if self.model not in model_names:
                print(f"Warning: Model '{self.model}' not found in Ollama.")
                print(f"Available models: {', '.join(model_names)}")
                print(f"You can pull the model using: ollama pull {self.model}")
                
        except requests.exceptions.RequestException as e:
            raise Exception(f"Cannot connect to Ollama server at {self.api_base}: {e}")

    @retry(
        wait=wait_random_exponential(min=1, max=10),
        stop=stop_after_attempt(3),
        retry=retry_if_exception_type(Exception),
    )
    def _chat_completion(self, messages: List[Dict]) -> str:
        """
        Send chat completion request to Ollama
        
        Parameters:
        -----------
        messages: List[Dict]
            List of message dictionaries with 'role' and 'content' keys
            
        Returns:
        --------
        str: Response content from the model
        """
        try:
            # Convert messages to Ollama format
            payload = {
                "model": self.model,
                "messages": messages,
                "stream": False
            }
            
            response = requests.post(
                f"{self.api_base}/api/chat", 
                json=payload, 
                timeout=120
            )
            
            if response.status_code != 200:
                raise Exception(f"Ollama API error: {response.status_code} - {response.text}")
            
            result = response.json()
            return result["message"]["content"]
            
        except requests.exceptions.Timeout:
            raise Exception("Ollama request timed out")
        except requests.exceptions.RequestException as e:
            raise Exception(f"Ollama request failed: {e}")
        except KeyError as e:
            raise Exception(f"Unexpected Ollama response format: {e}")

    def send_new_message(self, message: str, conversation_id: str = None) -> Tuple[str, str]:
        """
        Send a new message and start a new conversation
        
        Parameters:
        -----------
        message: str
            The message to send
        conversation_id: str, optional
            If provided, use this conversation ID
            
        Returns:
        --------
        Tuple[str, str]: (response_content, conversation_id)
        """
        if conversation_id is None:
            conversation_id = str(uuid1())
        
        # Create new conversation
        conversation = Conversation(conversation_id=conversation_id)
        self.conversation_dict[conversation_id] = conversation
        
        # Send message
        messages = [{"role": "user", "content": message}]
        
        msg_obj = Message(
            ask_id=str(uuid1()),
            ask={"role": "user", "content": message},
            request_start_timestamp=time.time(),
        )
        
        try:
            response_content = self._chat_completion(messages)
            
            msg_obj.request_end_timestamp = time.time()
            msg_obj.time_escaped = msg_obj.request_end_timestamp - msg_obj.request_start_timestamp
            msg_obj.answer = {"role": "assistant", "content": response_content}
            msg_obj.answer_id = str(uuid1())
            
            conversation.message_list.append(msg_obj)
            
            return response_content, conversation_id
            
        except Exception as e:
            logger.error(f"Error in send_new_message: {e}")
            raise e

    def send_message(self, message: str, conversation_id: str) -> str:
        """
        Send a message in an existing conversation
        
        Parameters:
        -----------
        message: str
            The message to send
        conversation_id: str
            The conversation ID to continue
            
        Returns:
        --------
        str: Response content from the model
        """
        if conversation_id not in self.conversation_dict:
            # If conversation doesn't exist, create new one
            response, _ = self.send_new_message(message, conversation_id)
            return response
        
        conversation = self.conversation_dict[conversation_id]
        
        # Build message history
        messages = []
        
        # Add recent messages from history (limited by history_length)
        recent_messages = conversation.message_list[-self.history_length:]
        for msg in recent_messages:
            if msg.ask:
                messages.append(msg.ask)
            if msg.answer:
                messages.append(msg.answer)
        
        # Add current message
        messages.append({"role": "user", "content": message})
        
        msg_obj = Message(
            ask_id=str(uuid1()),
            ask={"role": "user", "content": message},
            request_start_timestamp=time.time(),
        )
        
        try:
            response_content = self._chat_completion(messages)
            
            msg_obj.request_end_timestamp = time.time()
            msg_obj.time_escaped = msg_obj.request_end_timestamp - msg_obj.request_start_timestamp
            msg_obj.answer = {"role": "assistant", "content": response_content}
            msg_obj.answer_id = str(uuid1())
            
            conversation.message_list.append(msg_obj)
            
            return response_content
            
        except Exception as e:
            logger.error(f"Error in send_message: {e}")
            raise e

    def get_conversation_history(self, conversation_id: str = None) -> List[Dict]:
        """
        Get conversation history
        
        Parameters:
        -----------
        conversation_id: str, optional
            The conversation ID. If None, returns None
            
        Returns:
        --------
        List[Dict] or None: List of messages or None if conversation not found
        """
        if conversation_id is None or conversation_id not in self.conversation_dict:
            return None
        
        conversation = self.conversation_dict[conversation_id]
        history = []
        
        for msg in conversation.message_list:
            if msg.ask:
                history.append(msg.ask)
            if msg.answer:
                history.append(msg.answer)
        
        return history

    def refresh(self):
        """
        Refresh the API connection - for Ollama, just test connection
        """
        self._test_connection()
        logger.info("Ollama connection refreshed")

    def get_model_list(self) -> List[str]:
        """
        Get list of available models from Ollama
        
        Returns:
        --------
        List[str]: List of available model names
        """
        try:
            response = requests.get(f"{self.api_base}/api/tags", timeout=10)
            if response.status_code == 200:
                models = response.json().get("models", [])
                return [model["name"] for model in models]
            else:
                return []
        except Exception as e:
            logger.error(f"Error getting model list: {e}")
            return []